from rpyc.utils.security.base import BaseDB
from rpyc.utils.lib import SocketChannel
from rpyc.utils.security.base import ProtocolError, ProtocolAuthError
import random
try:
    from hashlib import sha1
except ImportError:
    from sha import sha as sha1


#
# username to hashed-password database
#
class UserDB(BaseDB):
    __slots__ = ["users"]
    SALT_BYTES = 10
    
    def __init__(self, users, filename = None):
        BaseDB.__init__(self, filename)
        self.users = users
    def __iter__(self):
        return self.users.iteritems()
    def __len__(self):
        return len(self.users)
    def __getitem__(self, username):
        return self.users[username]
    def __delitem__(self, username):
        del self.users[username]
    def __contains__(self, username):
        return username in self.users

    @classmethod
    def _load(cls, data):
        users = {}
        for line in data.splitlines():
            username, digest, salt = line.split()
            users[username] = (digest.decode("hex"), salt.decode("hex"))
        return cls(users)

    def _dump(self):
        lines = []
        for username, (digest, salt) in self:
            lines.append("%s %s %s" % (username, digest.encode("hex"), digest, salt.encode("hex")))
        return "\n".join(lines) + "\n"
    
    @classmethod
    def digest_new(cls, username, password):
        salt = "".join(chr(random.randint(0, 255)) for i in range(cls.SALT_BYTES))
        return cls.digest(username, password, salt), salt

    @classmethod
    def digest(cls, username, password, salt):
        return sha1(salt + username + password).digest()
    
    def test_password(self, username, password):
        try:
            digest, salt = self[username]
        except KeyError:
            digest, salt = None, "methuselah"
        digest2 = self.digest(username, password, salt)
        return digest == digest2
    
    def set_password(self, username, password):
        self.users[username] = self.digest_new(username, password)
    
    def list_users(self):
        return self.users.keys()

#
# username-password authentication protocol
#
class UserProtocol(object):
    @classmethod
    def handshake_server(cls, sock, userdb):
        chan = SocketChannel(sock)
        chan.send("USERNAME")
        username = chan.recv()
        chan.send("PASSWORD")
        password = chan.recv()
        if userdb.test_password(username, password):
            chan.send("OK")
        else:
            chan.send("ERROR")
            raise ProtocolAuthError("wrong username or password")
        return sock, username

    @classmethod
    def handshake_client(cls, sock, username, password):
        chan = SocketChannel(sock)
        chan.expect("USERNAME")
        chan.send(username)
        chan.expect("PASSWORD")
        chan.send(password)
        succ = chan.recv()
        if succ != "OK":
            raise ProtocolAuthError("server rejected username or password")
        return sock


